package edu.sysu.demo.rabbitmq.serial;

import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessageProperties;
import org.springframework.amqp.support.converter.AbstractMessageConverter;
import org.springframework.amqp.support.converter.MessageConversionException;

public class KryoMessageConverter extends AbstractMessageConverter {

  public static final String CONTENT_TYPE = "application/x-kryo";

  private final KryoSerializer kryoSerializer;

  private final boolean threadSafe;

  private final Charset defaultCharset;

  public KryoMessageConverter() {
    this(true, StandardCharsets.UTF_8);
  }

  public KryoMessageConverter(boolean threadSafe, Charset defaultCharset) {
    this.threadSafe = threadSafe;
    this.defaultCharset = defaultCharset;
    if (this.threadSafe) {
      kryoSerializer = KryoSerializer.threadLocal();
    } else {
      kryoSerializer = KryoSerializer.simple();
    }
  }

  @Override
  protected Message createMessage(Object object, MessageProperties messageProperties) {
    messageProperties.setContentType(CONTENT_TYPE);
    if (messageProperties.getContentEncoding() == null) {
      messageProperties.setContentEncoding(defaultCharset.name());
    }
    return new Message(kryoSerializer.serialize(object), messageProperties);
  }

  @Override
  public Object fromMessage(Message message) {
    MessageProperties properties = message.getMessageProperties();
    if (properties != null && properties.getContentType() != null
        && properties.getContentType().contains("x-kryo")) {
      return kryoSerializer.deserialize(message.getBody());
    }
    throw new MessageConversionException("Converter not applicable to this message");
  }
}
